package com.milvus.videoSearch;

import com.milvus.featureExtractor.ImageFeatureExtractor;
import com.milvus.featureExtractor.VideoFrameFeatureExtractor;
import com.milvus.imagesearch.ImageSearcher;
import com.milvus.imagesearch.MilvusManager;
import org.nd4j.linalg.api.ndarray.INDArray;

import java.io.File;
import java.io.IOException;
import java.io.Serializable;

/**
 * 视频搜索
 */
public class Main implements Serializable {
    private static final String COLLECTION_NAME = "image_collection";
    private static String VIDEO_PATH = "E:\\temp\\zhaojinmai.mp4"; // 视频文件路径

    public static void main(String[] args) throws IOException {
        ImageFeatureExtractor extractor = new ImageFeatureExtractor();
        MilvusManager milvusManager = new MilvusManager();
        VideoFrameFeatureExtractor videoFrameFeatureExtractor = new VideoFrameFeatureExtractor();
        ImageSearcher searcher = new ImageSearcher();

        // 创建Milvus集合
        //milvusManager.createCollection(COLLECTION_NAME);
        //videoFrameFeatureExtractor.extractFeatures(COLLECTION_NAME, VIDEO_PATH);

        //


        System.out.println("插入完成");

        // 查询图像
        File queryImage = new File("E:\\temp\\zhaojinmai\\frame_244.png");
        INDArray queryFeatures = extractor.extractFeatures(queryImage);
        searcher.search(COLLECTION_NAME,queryFeatures);
    }
}
